from datetime import datetime, timedelta
from fastapi import HTTPException, Request
from starlette.middleware.base import BaseHTTPMiddleware
from collections import defaultdict
from typing import Dict, Optional
from fastapi.responses import JSONResponse
from app.utils.log_config import get_logger

# 获取速率限制日志记录器
logger = get_logger("rate_limit")

class RateLimiter:
    def __init__(self):
        self._requests: Dict[str, Dict[str, list]] = defaultdict(lambda: defaultdict(list))
        # 添加默认配置
        self.default_limit = 10  # 默认每个时间窗口允许的请求数
        self.default_window = 60  # 默认时间窗口(秒)
        # 添加特定接口的限制配置
        self.endpoint_limits = {
            # 示例: "endpoint_name": (max_requests, window)
            "upload_file": (3, 60),  # 上传文件限制更严格
            "send_msg": (5, 60),     # 消息发送限制
            "get_report": (20, 60)   # 报告获取限制
        }
        # 添加白名单
        self.whitelist_users = set()
        # 添加豁免路径
        self.exempt_paths = {
            "/docs",           # Swagger文档
            "/redoc",         # ReDoc文档
            "/openapi.json",  # OpenAPI schema
            "/metrics",       # 监控指标
            "/health"         # 健康检查
        }
        
    def _clean_old_requests(self, user_id: str, endpoint: str, window: int):
        """清理过期的请求记录"""
        now = datetime.now()
        self._requests[user_id][endpoint] = [
            req for req in self._requests[user_id][endpoint]
            if now - req[0] < timedelta(seconds=window)
        ]

    def is_allowed(self, user_id: str, endpoint: str) -> bool:
        """检查请求是否允许"""
        # 检查是否在白名单中
        if user_id in self.whitelist_users:
            return True
            
        # 获取限制配置
        max_requests, window = self.endpoint_limits.get(
            endpoint, (self.default_limit, self.default_window)
        )
        
        # 清理过期请求
        self._clean_old_requests(user_id, endpoint, window)
        
        # 检查请求数量
        if len(self._requests[user_id][endpoint]) >= max_requests:
            logger.warning(f"Rate limit exceeded for user {user_id} on endpoint {endpoint}")
            return False
            
        # 记录新请求
        self._requests[user_id][endpoint].append((datetime.now(), endpoint))
        return True

    def add_to_whitelist(self, user_id: str):
        """添加用户到白名单"""
        self.whitelist_users.add(user_id)

    def remove_from_whitelist(self, user_id: str):
        """从白名单移除用户"""
        self.whitelist_users.discard(user_id)

    def set_endpoint_limit(self, endpoint: str, max_requests: int, window: int):
        """设置特定接口的限制"""
        self.endpoint_limits[endpoint] = (max_requests, window)

# 创建全局实例
limiter = RateLimiter()

class RateLimitMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        # 检查是否是豁免路径
        if request.url.path in limiter.exempt_paths:
            return await call_next(request)
            
        try:
            # 获取用户信息
            user = request.scope.get("user")
            if not user:
                # 如果找不到用户信息，可以使用IP地址作为标识
                user_id = request.client.host
            else:
                user_id = str(user.id)
                
            # 获取endpoint名称
            endpoint = request.url.path.split("/")[-1]
            
            if not limiter.is_allowed(user_id, endpoint):
                max_requests, window = limiter.endpoint_limits.get(
                    endpoint, (limiter.default_limit, limiter.default_window)
                )
                return JSONResponse(
                    status_code=429,
                    content={
                        "detail": f"请求过于频繁，请在{window}秒后重试",
                        "max_requests": max_requests,
                        "window": window
                    }
                )
                
            response = await call_next(request)
            return response
            
        except Exception as e:
            logger.error(f"Rate limit middleware error: {str(e)}")
            return await call_next(request) 